// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::ops::{Deref, DerefMut};

use thiserror::Error;
use tracing::trace;

#[derive(Debug, Error)]
pub enum Error {
    #[error("Insufficient buffer capacity (required: {required}, capacity: {capacity})")]
    InsufficientBuffer { required: usize, capacity: usize },

    #[error("Frame magic number mismatch")]
    MagicNumberMismatch,

    #[error("Frame length exceeds limit (limit: {limit}, actual: {actual})")]
    FrameTooLarge { limit: usize, actual: usize },

    #[error("Frame size overflow")]
    FrameSizeOverflow,
}

#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct FrameHeader {
    magic: u64,
    length: usize,
}

impl FrameHeader {
    const MAGIC: u64 = 0x78475055_5F464D45; // 'xGPU_FME' in ASCII
    const SIZE: usize = size_of::<Self>();

    const MAGIC_START: usize = 0;
    const MAGIC_END: usize = size_of::<u64>();
    const LEN_START: usize = FrameHeader::SIZE - size_of::<usize>();
    const LEN_END: usize = FrameHeader::SIZE;

    #[inline(always)]
    fn write_to(&self, buf: &mut [u8]) {
        debug_assert!(buf.len() >= Self::SIZE);
        buf[Self::MAGIC_START..Self::MAGIC_END].copy_from_slice(&self.magic.to_ne_bytes());
        buf[Self::LEN_START..Self::LEN_END].copy_from_slice(&self.length.to_ne_bytes());
    }

    #[inline(always)]
    fn read_from(buf: &[u8]) -> Self {
        debug_assert!(buf.len() >= Self::SIZE);
        Self {
            magic: u64::from_ne_bytes(buf[Self::MAGIC_START..Self::MAGIC_END].try_into().unwrap()),
            length: usize::from_ne_bytes(buf[Self::LEN_START..Self::LEN_END].try_into().unwrap()),
        }
    }
}

#[derive(Debug)]
pub struct Frame<B> {
    payload: B,
}

impl<B: AsRef<[u8]>> Deref for Frame<B> {
    type Target = [u8];

    #[inline]
    fn deref(&self) -> &Self::Target {
        self.payload.as_ref()
    }
}

impl<B: AsRef<[u8]>> Frame<B> {
    #[inline]
    pub fn frame_len(&self) -> usize {
        FrameHeader::SIZE + self.payload.as_ref().len()
    }
}

#[derive(Debug)]
pub struct FrameBuffer<B> {
    buf: B,
    limit: usize,
}

impl<B: AsRef<[u8]> + AsMut<[u8]>> FrameBuffer<B> {
    pub fn finalize(mut self, payload_len: usize) -> Result<usize, Error> {
        let frame_len = FrameHeader::SIZE
            .checked_add(payload_len)
            .ok_or(Error::FrameSizeOverflow)?;
        if frame_len > self.limit {
            return Err(Error::FrameTooLarge {
                limit: self.limit,
                actual: frame_len,
            });
        }

        let buf_len = self.buf.as_ref().len();
        if frame_len > buf_len {
            return Err(Error::InsufficientBuffer {
                required: frame_len,
                capacity: buf_len,
            });
        }

        let (header_buf, _) = self.buf.as_mut().split_at_mut(FrameHeader::SIZE);

        let header = FrameHeader {
            magic: FrameHeader::MAGIC,
            length: payload_len,
        };
        header.write_to(header_buf);

        trace!(
            "[Frame] Encoded {} bytes (payload {} bytes)",
            frame_len,
            payload_len
        );
        Ok(frame_len)
    }
}

impl<B: AsRef<[u8]>> Deref for FrameBuffer<B> {
    type Target = [u8];

    #[inline]
    fn deref(&self) -> &Self::Target {
        debug_assert!(self.buf.as_ref().len() >= FrameHeader::SIZE);
        &self.buf.as_ref()[FrameHeader::SIZE..]
    }
}

impl<B: AsRef<[u8]> + AsMut<[u8]>> DerefMut for FrameBuffer<B> {
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        debug_assert!(self.buf.as_ref().len() >= FrameHeader::SIZE);
        &mut self.buf.as_mut()[FrameHeader::SIZE..]
    }
}

#[derive(Debug, Clone, Copy)]
pub struct Framer {
    limit: usize,
}

impl Framer {
    #[inline]
    pub fn new(limit: usize) -> Self {
        assert!(limit >= FrameHeader::SIZE);
        Self { limit }
    }
}

impl Default for Framer {
    #[inline]
    fn default() -> Self {
        Self::new(usize::MAX)
    }
}

impl Framer {
    pub fn encode_frame<'a>(&self, buf: &'a mut [u8]) -> Result<FrameBuffer<&'a mut [u8]>, Error> {
        let buf_len = buf.len();

        trace!("[Frame] Encoding frame, buf_len={}...", buf_len);

        if buf_len < FrameHeader::SIZE {
            return Err(Error::InsufficientBuffer {
                required: FrameHeader::SIZE,
                capacity: buf_len,
            });
        }

        let buffer = FrameBuffer {
            buf,
            limit: self.limit,
        };
        Ok(buffer)
    }

    pub fn decode_frame<'a>(&self, buf: &'a [u8]) -> Result<Option<Frame<&'a [u8]>>, Error> {
        let buf_len = buf.len();

        trace!("[Frame] Decoding frame, buf_len={}...", buf_len);

        /* Skip if no enough data for a full header */
        if buf_len < FrameHeader::SIZE {
            return Ok(None);
        }

        /* Parse frame header */
        let header = FrameHeader::read_from(buf);

        /* Check frame magic number */
        if header.magic != FrameHeader::MAGIC {
            return Err(Error::MagicNumberMismatch);
        }

        /* Check frame size */
        let frame_len = FrameHeader::SIZE
            .checked_add(header.length)
            .ok_or(Error::FrameSizeOverflow)?;
        if frame_len > self.limit {
            return Err(Error::FrameTooLarge {
                limit: self.limit,
                actual: frame_len,
            });
        }

        /* Skip if no enough data for the whole frame */
        if buf_len < frame_len {
            return Ok(None);
        }

        /* Get frame payload */
        let payload = &buf[FrameHeader::SIZE..frame_len];

        trace!(
            "[Frame] Decoded {} bytes (payload {} bytes)",
            frame_len,
            header.length
        );
        Ok(Some(Frame { payload }))
    }
}
